using System;
using System.Collections.Generic;
using System.Text;
using DotNetMatrix;
using Blimpbots.Recognition.Algorithms;


namespace Blimpbots.Recognition
{
    public class RansacAffine
    {
        public static long Choose(int n, int r)
        {
            // Implements C(n,r)
            //  = n!/(r!(n-r)!)

            long numerator = 1;
            for (int i = n; i > (n - r); i--)
                numerator *= i;

            long denominator = 1;
            for (int i = 1; i <= r; i++)
                denominator *= i;

            return (numerator / denominator);

        }


        /// <summary>
        /// From the shuffled list, draws matches in order, verifying that no two
        /// drawn refer to the same model keypoint.
        /// </summary>
        private static bool GetUniqueSet(ShuffledList<KeypointMatch> sl, int desiredN, out List<KeypointMatch> set)
        {
            set = new List<KeypointMatch>(desiredN);
            for (int i = 0; i < sl.Count; i++)
            {
                bool collide = false;
                foreach (KeypointMatch k in set)
                {
                    if (k.modelIdx == sl[i].modelIdx) { collide = true; break; }
                    if (k.dataIdx == sl[i].dataIdx) { collide = true; break; }
                }
                if (!collide)
                {
                    set.Add(sl[i]);
                    if (set.Count == desiredN) return true;
                }
            }
            return false; // ran out of items before hitting desired count
        }

        static Random rand = new Random();

        #region Old Spatial Selector (May be broken)
        //private static bool GetSpatialSet(ShuffledList<KeypointMatch> sl, int desiredN, out KeypointMatch[] set)
        //{
        //    set = null;

        //    // We need 1 for pivot plus N to find a set
        //    if (sl.Count <= desiredN) return false;

        //    // Choose one keypoint to select around
        //    KeypointMatch pivot = sl[0];
        //    double TwoPi = 2 * Math.PI;
        //    double randAng = TwoPi * rand.NextDouble(); 
        //    double angleSpan = TwoPi / desiredN;

        //    Point2D pivotPt = pivot.dataCoord;
        //    int goodPtsFound = 0;

        //    set = new KeypointMatch[desiredN];
        //    for (int i = 1; i < sl.Count; i++)
        //    {
        //        KeypointMatch prospect = sl[i];
        //        bool collide = false;
        //        foreach (KeypointMatch k in set)
        //        {
        //            if (k == null) continue; // Skip unset items
        //            if (k.modelIdx == prospect.modelIdx) { collide = true; break; }
        //            if (k.dataIdx == prospect.dataIdx) { collide = true; break; }
        //        }
        //        if (!collide)
        //        {
        //            double prospectAng = Math.Atan2(prospect.dataCoord.y - pivotPt.y,
        //                                            prospect.dataCoord.x - pivotPt.x);
        //            if (prospectAng < 0) prospectAng += TwoPi;
        //            prospectAng -= randAng;
        //            if (prospectAng < 0) prospectAng += TwoPi; // pull into 0...2Pi
        //            int idx = (int)(prospectAng / angleSpan);
        //            if (set[idx] == null)
        //            {
        //                goodPtsFound++; 
        //                set[idx] = prospect;
        //            }
        //            if (goodPtsFound == desiredN) return true;
        //        }
        //    }
        //    return false; // ran out of items before hitting desired count
        //}
        #endregion

        private static bool GetSpatialSet(ShuffledList<KeypointMatch> sl, int desiredN, out List<KeypointMatch> set)
        {
            double bestDist = 0;
            set = new List<KeypointMatch>(desiredN);
            for (int i = 1; i < sl.Count; i++)
            {
                KeypointMatch prospect = sl[i];
                bool collide = false;
                foreach (KeypointMatch k in set)
                {
                    if (k.modelIdx == prospect.modelIdx) { collide = true; break; }
                    if (k.dataIdx == prospect.dataIdx) { collide = true; break; }
                }
                if (!collide)
                    if (set.Count < desiredN - 1)
                        set.Add(prospect);
                    else if( set.Count == desiredN - 1)
                    {
                        set.Add(prospect);
                        for(int j = 0; j < desiredN - 1; j++)
                            bestDist +=  Math.Pow((prospect.modelCoord.x - set[j].modelCoord.x),2)
                                       + Math.Pow((prospect.modelCoord.y - set[j].modelCoord.y),2);
                    }
                    else
                    {
                        double dist = 0; // Sum up the distance to the existing set
                        for(int j = 0; j < desiredN - 1; j++)
                            dist +=  Math.Pow((prospect.modelCoord.x - set[j].modelCoord.x),2)
                                   + Math.Pow((prospect.modelCoord.y - set[j].modelCoord.y),2);

                        // bigger distance is better
                        if (dist > bestDist)
                        {
                            bestDist = dist;
                            set[desiredN - 1] = prospect;
                        }

                        // Randomly quit early to reduce determinism
                        if (rand.Next(sl.Count) > sl.Count / 2) { return true; }
                    }

            } // end foreach in shuffled list

            if (set.Count == desiredN) return true;
            else
               return false; // ran out of items before hitting desired count
        }


        /// <summary>
        /// Given an affine transform, try to fit as many examples as possible, and refine
        /// the transform afterward.
        /// </summary>
        /// <param name="matches"></param>
        /// <param name="t"></param>
        /// <returns></returns>
        static public AffineModel TryRefit(List<KeypointMatch> matches,
                                             AffineTransform t, double errThresh)
        {
            // not enough to begin with ...
            if (matches.Count < 4) return null;
            
            AffineModel m = new AffineModel(matches, t, errThresh);
            if (m.Inliers.Count < 4) return null;
            m.TryRefine();
            return m;
        }


        // May return null if no good models can be found
        static public AffineModel Run(List<KeypointMatch> matches, int n, 
                                        double goodScore, double errThresh, int maxRuns)
        {
            ShuffledList<KeypointMatch> sl = new ShuffledList<KeypointMatch>(matches);

            AffineModel bestModel = null;

            for (int run = 0; run < maxRuns; run++)
            {
                sl.Shuffle(1); // Randomize the list

                //List<KeypointMatch> set;
                //if (!GetSpatialSet(sl, n, out set))
                //    break; // Give up if there ever aren't enough items

                List<KeypointMatch> set;
                if (!GetUniqueSet(sl, n, out set))
                    break; // Give up if there ever aren't enough items

                AffineModel m = new AffineModel(set, errThresh);

                // Add all of the other matches that agree
                for (int i = n; i < sl.Count; i++)
                    m.TryAdd(sl[i]);

                // If we added some, try to refine this one
                if(m.Inliers.Count > n)
                    m.TryRefine();

                if (bestModel == null)
                    bestModel = m;
                else
                {
                    if (bestModel.Score > m.Score)
                        bestModel = m;
                }

                // Short-circuit if we've found a good enough fit
                if (bestModel.Score < goodScore)
                    break;
            }

            return bestModel;
        }
    }

    public class AffineModel
    {
        private AffineTransform _t = null;
        private IList<KeypointMatch> _matches;
        private double _errThresh;

        private double _sumError = 0;
        private double _errPenalty;

        // Models above this size are not penalized beyond
        // their reprojection error.
        private const int goodSize = 10;

        public GeneralMatrix H {
            get { return _t.H; }
        }

        public double Rotation
        {
            get { return _t.Rotation; }
        }

        public AffineTransform Transform { get { return _t; } }

        public IList<KeypointMatch> Inliers
        {
            get { return _matches; }
        }

        public double Score
        {
            get
            {
                return (1.0 / _matches.Count) * (_t.Skew); 
            }
        }

        public double Scale
        {
            get
            {
                return _t.Scale;
            }
        }

        // Builds an AffineModel given a transform, including only those matches
        // which are consistent.
        public AffineModel(IList<KeypointMatch> possibleMatches, AffineTransform t,
                             double errThresh )
        {
            _matches = new List<KeypointMatch>();

            _errThresh = errThresh;
            _t = t;

            Point2D err = AvgDistance(possibleMatches);
            _t.H.Array[0][2] -= err.x;
            _t.H.Array[1][2] -= err.y;
                       
            foreach (KeypointMatch m in possibleMatches)
                TryAdd(m);
        }


        public AffineModel(IList<KeypointMatch> matches, double errThresh)
        {
            // Remember all the model matches
            _errThresh = errThresh;

            _matches = matches;

            GeneralMatrix x1, x2; 
            BuildPairs(_matches, _matches.Count, out x1, out x2);

            _errPenalty = EstimateMaxError(x2);
            _t = new AffineTransform(x1, x2);

            _sumError = _t.SelfError;
        }

        // Estimates the maximum real-world error for an affine model
        // that maps to these points
        private double EstimateMaxError(GeneralMatrix x2)
        {
            if (x2.ColumnDimension > goodSize)
                return 0;

            if (x2.ColumnDimension == 0) 
                throw new Exception("Cannot estimate error on zero-sized object.");

            int n = x2.ColumnDimension;

            double distSum = 0;
            for (int i = 0; i < n; i++)
                for (int j = 0; j < n; j++)
                {
                    if (j == i) continue;
                    // Sum up the distance between pairs
                    distSum += Math.Sqrt(Math.Pow(x2.Array[0][i] - x2.Array[0][j], 2) +
                               Math.Pow(x2.Array[1][i] - x2.Array[1][j], 2));
                }

            return distSum / (n*(n-1)); // number of pairs is n*(n-1)

        }

        private void BuildPairs(IEnumerable<KeypointMatch> matches, int count, out GeneralMatrix mx1, out GeneralMatrix mx2)
        {
            #region Build matrix of homogeneous coordinates

            double[][] x1 = new double[3][];
            x1[0] = new double[count];
            x1[1] = new double[count];
            x1[2] = new double[count];

            double[][] x2 = new double[3][];
            x2[0] = new double[count];
            x2[1] = new double[count];
            x2[2] = new double[count];

            IEnumerator<KeypointMatch> kE = matches.GetEnumerator();

            for (int i = 0; i < count; i++)
            {
                kE.MoveNext();
                KeypointMatch curMatch = kE.Current;

                x1[0][i] = curMatch.modelCoord.x;
                x1[1][i] = curMatch.modelCoord.y;
                x1[2][i] = 1.0;

                x2[0][i] = curMatch.dataCoord.x;
                x2[1][i] = curMatch.dataCoord.y;
                x2[2][i] = 1.0;
            }

            #endregion

            mx1 = new GeneralMatrix(x1);
            mx2 = new GeneralMatrix(x2);

        }

        /// <summary>
        /// Finds the average offset between the projected and actual
        /// position for the points given.
        /// </summary>
        public Point2D AvgDistance(IList<KeypointMatch> matches)
        {
            double xerr_sum = 0, yerr_sum = 0;
            GeneralMatrix x1, x2;
            int n = matches.Count;
            BuildPairs(matches, n, out x1, out x2);
            GeneralMatrix x2_est = _t.H * x1;
            for (int i = 0; i < n; i++)
            {
                xerr_sum += x2_est.Array[0][i] - x2.Array[0][i];
                yerr_sum += x2_est.Array[1][i] - x2.Array[1][i];
            }
            return new Point2D(xerr_sum / n, yerr_sum / n);
        }

        // Add a match to the model if it's within the
        // allowable error bounds AND if the same
        // model keypoint hasn't already been used
        public bool TryAdd(KeypointMatch m)
        {
            // Check for dupe model keypoints
            foreach (KeypointMatch existingMatch in _matches)
            {
                if (m.modelIdx == existingMatch.modelIdx) return false;
                if (m.dataIdx == existingMatch.dataIdx) return false;               
            }

            GeneralMatrix x1, x2;
            BuildPairs(new KeypointMatch[]{m}, 1, out x1, out x2);
            double err = _t.CalcError(x1, x2);
            if (err < _errThresh)
            {
                _sumError += err;
                _matches.Add(m);
                return true;
            }
            else
                return false;
        }

        // Check if computing a new AffineTransform
        // using all of the model's full set gives a better
        // fit than the initially computed transform
        public void TryRefine()
        {
            GeneralMatrix x1, x2;
            BuildPairs(_matches, _matches.Count, out x1, out x2);
            AffineTransform _candidateT = new AffineTransform(x1, x2);

            if (_candidateT.SelfError < _sumError)
            {
                _t = _candidateT;
                _sumError = _candidateT.SelfError;
            }
        }

    }

    public class AffineTransform
    {
        public GeneralMatrix H { get { return _h; } }
        public double Rotation { get { return _rotation; } }

        public double SelfError { get { return _error; } }

        public Point2D Offset
        {
            get
            {
                return new Point2D(_h.Array[0][2], _h.Array[1][2]);
            }

        }

        GeneralMatrix _h;
        double _error; // public error after best-fit

        public double CalcError(GeneralMatrix x1, GeneralMatrix x2)
        {
            GeneralMatrix x2_est = _h * x1;
            int n = x1.ColumnDimension;
            double err = 0;
            for (int i = 0; i < n; i++)
            {
                double xerr = x2_est.Array[0][i] - x2.Array[0][i];
                double yerr = x2_est.Array[1][i] - x2.Array[1][i];
                err += (xerr * xerr) + (yerr * yerr);
            }
            return err;
        }

        private bool _decomposed = false;
        private double _skew, _scale, _rotation;

        // http://www.cs.umd.edu/~djacobs/CMSC426/LinearAlgebraSVD.pdf

        private void Decompose()
        {
            GeneralMatrix affine2x2 = new GeneralMatrix(2, 2);
            affine2x2.SetMatrix(0, 1, 0, 1, H);
            SingularValueDecomposition svd = affine2x2.SVD();

            _skew = svd.SingularValues[0] / svd.SingularValues[1];
            _scale = svd.SingularValues[1];

            GeneralMatrix u = svd.GetU();
            
            _rotation = Math.Atan2(u.Array[0][0], u.Array[0][1]);

            _decomposed = true;
        }

        #region Decomposition Accessors

        public double Scale
        {
            get
            {
                if (!_decomposed) Decompose();
                return _scale;
            }
        }

        public double Skew
        {
            get
            {
                if (!_decomposed) Decompose(); 
                return _skew;
            }
        }

        #endregion

        public AffineTransform(GeneralMatrix x1, GeneralMatrix x2)
        {
            // Ported from  vgg_Haffine_from_x_MLE(xs1,xs2) MATLAB code
            // MATLAB author: Andrew Zisserman

            int numpoints = x1.ColumnDimension;

            # region Input Validation

            if (numpoints != x2.ColumnDimension)
                throw new Exception("Input columns don't match!");
            else if (numpoints < 2)
                throw new Exception("Must specify at least 2 point matches.");

            if (x1.RowDimension != x2.RowDimension || x1.RowDimension != 3)
                throw new Exception("Row dimension must be 3.");

            #endregion

            double m1x = x1.RowMean(0);
            double m2x = x2.RowMean(0);
            double m1y = x1.RowMean(1);
            double m2y = x2.RowMean(1);

            double sumsqerrx = 0, sumsqerry = 0;
            for (int i = 0; i < numpoints; i++)
            {
                sumsqerrx += Math.Pow(x1.Array[0][i] - m1x, 2);
                sumsqerry += Math.Pow(x1.Array[1][i] - m1y, 2);
            }

            // Compute maximum standard deviation
            // StD = Sqrt(MSE) (normalized by N-1)
            double maxStd =
                Math.Max(Math.Sqrt(sumsqerrx / (numpoints - 1)),
                Math.Sqrt(sumsqerry / (numpoints - 1)));

            // C is initially the C1 matrix (from the MATLAB code)
            GeneralMatrix C1 = new GeneralMatrix(new double[3][] 
                { new double[3] { 1/maxStd,   0,   -m1x/maxStd } ,
                  new double[3] { 0,   1/maxStd,   -m1y/maxStd } ,
                  new double[3] { 0,          0,             1 } ,
                });

            GeneralMatrix Cx1 = C1 * x1;

            // Convert C1 to the C2 matrix
            GeneralMatrix C2 = (GeneralMatrix)C1.Clone();
            C2.Array[0][2] = -m2x / maxStd;
            C2.Array[1][2] = -m2y / maxStd;

            GeneralMatrix Cx2 = C2 * x2;

            // Normalize and transpose.  
            // MATLAB: 
            //   xs1nh = vgg_get_nonhomg(xs1);
            //   xs2nh = vgg_get_nonhomg(xs2);
            //   A = [xs1nh;xs2nh]';

            GeneralMatrix A = new GeneralMatrix(numpoints, 4);
            for (int i = 0; i < numpoints; i++)
            {
                A.Array[i][0] = Cx1.Array[0][i] / Cx1.Array[2][i];
                A.Array[i][1] = Cx1.Array[1][i] / Cx1.Array[2][i];
                A.Array[i][2] = Cx2.Array[0][i] / Cx2.Array[2][i];
                A.Array[i][3] = Cx2.Array[1][i] / Cx2.Array[2][i];
            }

            // Perform the hard work
            SingularValueDecomposition svd = A.SVD();

            //GeneralMatrix u = svd.GetU();
            //GeneralMatrix s = svd.S;
            GeneralMatrix v = svd.GetV();

            double epsilon = 2.2204E-016 * 1000;
            int nullDimension = 0;
            for (int i = 0; i < svd.SingularValues.Length; i++)
                if (svd.SingularValues[i] < (epsilon * svd.SingularValues[1]))
                    nullDimension += 1;

            if (nullDimension > 2) throw new Exception("Nullspace is too large.");

            GeneralMatrix B = new GeneralMatrix(2, 2);
            GeneralMatrix C = new GeneralMatrix(2, 2);

            for (int i = 0; i < 2; i++)
                for (int j = 0; j < 2; j++)
                {
                    B.Array[i][j] = v.Array[i][j];
                    C.Array[i][j] = v.Array[i + 2][j];
                }

            GeneralMatrix UL = C * B.Inverse();

            _h = new GeneralMatrix(new double[3][] 
                { new double[3] { UL.Array[0][0], UL.Array[0][1],   0 } ,
                  new double[3] { UL.Array[1][0], UL.Array[1][1],   0 } ,
                  new double[3] {             0,             0,   1 } ,
                });

            _h = C2.Inverse() * _h * C1;            // H = inv(C2) * H * C1;
            _h = _h.Multiply(1.0 / _h.Array[2][2]);  // H = H / H(3, 3);

            _error = CalcError(x1, x2); // Figure out how good the fit was

        }

    }

}
